import numpy as np
import akshare as ak
import pandas as pd
from sqlalchemy import create_engine, text, exc
from datetime import datetime, timedelta
import time
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('stock_data_single_bfq.log'),
        logging.StreamHandler()
    ]
)

# 新数据库配置
POSTGRES_SYSTEM_URL = "postgresql://leizhen:751982leizhen@localhost:5432/postgres"
DATABASE_URL = "postgresql://leizhen:751982leizhen@localhost/stock_db_single_bfq"
engine = create_engine(DATABASE_URL, pool_size=20, max_overflow=10)

def init_database():
    """初始化股票数据库表结构（单表模式）"""
    temp_engine = create_engine(POSTGRES_SYSTEM_URL)
    try:
        test_engine = create_engine(DATABASE_URL)
        with test_engine.connect():
            pass
        logging.info("数据库已存在，继续初始化")
    except exc.OperationalError:
        logging.info("创建新数据库stock_db_single_bfq...")
        try:
            with temp_engine.connect() as conn:
                conn.execute(text("COMMIT"))
                conn.execute(text("CREATE DATABASE stock_db_single_bfq"))
                logging.info("数据库创建成功")
        except Exception as e:
            logging.error(f"创建数据库失败: {str(e)}")
            raise SystemExit("数据库创建失败")

    # 创建股票基本信息表
    with engine.connect() as conn:
        conn.execute(text("""
            CREATE TABLE IF NOT EXISTS stocks (
                symbol VARCHAR(10) PRIMARY KEY,
                name VARCHAR(100) NOT NULL,
                industry VARCHAR(50),
                listing_date DATE NOT NULL,
                last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """))
        conn.commit()

def get_all_stocks():
    """获取全量A股列表（与原代码相同）"""
    try:
        df = ak.stock_info_a_code_name()
        return df[['code', 'name']].rename(columns={'code': 'symbol', 'name': 'name'})
    except Exception as e:
        logging.error(f"获取股票列表失败: {str(e)}")
        return pd.DataFrame()

def get_stock_history(symbol, start_date):
    """获取个股历史行情数据（与原代码相同）"""
    max_retries = 3
    for attempt in range(max_retries):
        try:
            df = ak.stock_zh_a_hist(
                symbol=symbol,
                period="daily",
                start_date=start_date.strftime("%Y%m%d"),
                end_date=datetime.now().strftime("%Y%m%d"),
                adjust="")
            return df.rename(columns={
                '日期': 'trade_date',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'volume',
                '股票代码': 'symbol',
                '成交额': 'amount',
                '振幅': 'price_amplitude',
                '涨跌幅': 'change_pct',
                '涨跌额': 'change_amount',
                '换手率': 'turnover_rate'
            })[['symbol', 'trade_date', 'open', 'close', 'high', 'low', 'volume', 
               'amount', 'price_amplitude', 'change_pct', 'change_amount', 'turnover_rate']]
        except Exception as e:
            if attempt == max_retries - 1:
                logging.warning(f"获取{symbol}数据失败: {str(e)}")
                return pd.DataFrame()
            time.sleep(2 ** attempt)
    return pd.DataFrame()

def create_individual_table(conn, symbol):
    """为单个股票创建独立数据表"""
    conn.execute(text(f"""
        CREATE TABLE IF NOT EXISTS stock_{symbol} (
            id SERIAL PRIMARY KEY,
            trade_date DATE UNIQUE,
            open NUMERIC(10,2),
            close NUMERIC(10,2),
            high NUMERIC(10,2),
            low NUMERIC(10,2),
            volume BIGINT,
            amount NUMERIC(15,2),
            price_amplitude NUMERIC(5,2),
            change_pct NUMERIC(5,2),
            change_amount NUMERIC(10,2),
            turnover_rate NUMERIC(5,2)
        )
    """))
    conn.commit()

def main():
    init_database()
    
    stocks_df = get_all_stocks()
    if stocks_df.empty:
        logging.error("未获取到股票列表")
        return

    stocks_df.to_sql('stocks', engine, if_exists='replace', index=False)
    
    start_date = datetime.now() - timedelta(days=30*365)
    
    success = 0
    total = len(stocks_df)
    for idx, row in stocks_df.iterrows():
        symbol = row['symbol']
        name = row['name']
        
        logging.info(f"正在处理 [{idx+1}/{total}] {symbol}-{name}")
        
        history = get_stock_history(symbol, start_date)
        if not history.empty:
            # 移除symbol列并清洗数据
            history = history[['trade_date', 'open', 'close', 'high', 'low', 'volume', 
                             'amount', 'price_amplitude', 'change_pct', 'change_amount', 'turnover_rate']].copy()
            history.drop_duplicates(subset=['trade_date'], inplace=True)

            try:
                history['trade_date'] = pd.to_datetime(history['trade_date'], format='%Y-%m-%d').dt.date
                # 转换百分比值为小数
                percent_cols = ['change_pct', 'turnover_rate', 'price_amplitude']
                for col in percent_cols:
                    history[col] = history[col].astype(float) / 100.0
            except Exception as e:
                logging.error(f"日期格式转换失败: {str(e)}")
                continue

            with engine.connect() as db_conn:
                # 创建个股独立表
                create_individual_table(db_conn, symbol)
                
                # 检查现有数据
                existing = pd.read_sql(f"SELECT trade_date FROM stock_{symbol}", db_conn)
                existing_dates = set(existing['trade_date']) if not existing.empty else set()
                
                # 过滤新数据
                new_records = history[~history['trade_date'].isin(existing_dates)]
                
                if not new_records.empty:
                    # 数据清洗
                    new_records = new_records.replace([np.inf, -np.inf], np.nan).dropna()
                    new_records = new_records.round({
                        'open': 2, 'close': 2, 'high': 2, 'low': 2,
                        'amount': 2, 'price_amplitude': 2, 
                        'change_pct': 2, 'change_amount': 2, 'turnover_rate': 2
                    })
                    
                    # 插入数据（修复列名问题）
                    try:
                        # 明确指定要插入的列（移除symbol列）
                        insert_cols = ['trade_date', 'open', 'close', 'high', 'low', 
                                     'volume', 'amount', 'price_amplitude', 
                                     'change_pct', 'change_amount', 'turnover_rate']
                        
                        # 使用更安全的参数化插入
                        db_conn.execute(
                            text(f"""
                                INSERT INTO stock_{symbol} 
                                (trade_date, open, close, high, low, volume, 
                                 amount, price_amplitude, change_pct, change_amount, turnover_rate)
                                VALUES 
                                (:trade_date, :open, :close, :high, :low, :volume,
                                 :amount, :price_amplitude, :change_pct, :change_amount, :turnover_rate)
                                ON CONFLICT (trade_date) DO NOTHING
                            """),
                            [row.to_dict() for _, row in new_records[insert_cols].iterrows()]
                        )
                        db_conn.commit()
                        logging.info(f"成功插入{len(new_records)}条新记录到stock_{symbol}")
                    except Exception as e:
                        logging.error(f"插入数据失败: {str(e)}")
                        logging.debug(f"问题数据样本: {new_records.head(1).to_dict()}")
                    
                    success += 1
                    time.sleep(1)

    logging.info(f"数据采集完成，共处理{success}/{total}只股票")

if __name__ == "__main__":
    main()
